/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <folly/Random.h>
#include <gtest/gtest.h>

#include "cachelib/common/AccessTracker.h"

namespace facebook {
namespace cachelib {

namespace detail {
// A not thread-safe implementation of a ticker powered
// by an underlying tick stamp.
class NumberTicker : public Ticker {
 public:
  void setTicks(uint32_t ticks) { ticks_ = ticks; }
  void advanceTicks() { ticks_++; }
  virtual uint32_t getCurrentTick() override { return ticks_; }

 private:
  uint32_t ticks_{0};
};

} // namespace detail

class AccessTrackerTest : public ::testing::TestWithParam<bool> {
 protected:
  // Assert two feature vectors generated by tracker are equal.
  void assertVecEq(const std::vector<double>& actual,
                   const std::vector<double>& expected) {
    ASSERT_EQ(actual.size(), expected.size());
    for (size_t i = 0; i < actual.size(); i++) {
      ASSERT_EQ(actual[i], sanitizeCount(expected[i]));
    }
  }

  // Convert count to BF result.
  double sanitizeCount(double count) {
    return GetParam() ? count : (count > 0 ? 1 : 0);
  }

 protected:
  void advanceTicks() { ticker_->advanceTicks(); }

  // Start from the first tick of the bucket so that test cases
  // are not disturbed by randomization.
  void initializeTicks(size_t ticksPerBucket,
                       size_t numBuckets,
                       size_t numTicksPerBucket) {
    auto ticks = initialTicks - initialTicks % ticksPerBucket;

    // If the initial bucket is numBuckets - 1, advance it (to bucket 0).
    //
    // When trying to update the mostRecnet bucket, if the current bucket
    // index is 1 bucket before the mostRecent bucket, we would assume
    // that a race happened and the current time is actually at the next
    // bucket. We we would use the next bucket.
    // Historically we initialize the mostRecent bucket to 0. This brings
    // an issue if the initial time starts at bucket numBuckets - 1.
    // For the first hour, it would think it is very close to the next
    // bucket for up to an entire hour.
    // Making this change in the unit test elimiates the gap while
    // we consider a fix for the history issue.
    if ((ticks / numTicksPerBucket) % numBuckets == numBuckets - 1) {
      ticks += numTicksPerBucket;
    }

    ticker_->setTicks(ticks);
  }

  std::shared_ptr<detail::NumberTicker> ticker_ =
      std::make_shared<detail::NumberTicker>();

 private:
  // Random initial tick number.
  size_t initialTicks{folly::Random::rand32()};
};
INSTANTIATE_TEST_CASE_P(AccessTrackerTest, AccessTrackerTest, testing::Bool());

TEST_P(AccessTrackerTest, simpleTestCase) {
  auto config = AccessTracker::Config();
  config.numBuckets = 3;
  config.useCounts = GetParam();
  // Moves to a new bucket every two accesses.
  config.numTicksPerBucket = 2;
  config.ticker = ticker_;
  initializeTicks(
      config.numTicksPerBucket, config.numBuckets, config.numTicksPerBucket);
  auto tracker = AccessTracker(std::move(config));

  // The count is taken after the access tracking vector is returned.
  // After each assertion, the line below shows the status of the buckest, in
  // natural index order. The expected output order is rotated to start from
  // the most recent admission.
  folly::StringPiece key0 = "key0";
  // Bucket 0.
  assertVecEq(tracker.recordAndPopulateAccessFeatures(key0), {0, 0, 0});
  // key0: {1, 0, 0}
  advanceTicks();
  assertVecEq(tracker.recordAndPopulateAccessFeatures(key0), {1, 0, 0});
  // key0: {2, 0, 0}
  advanceTicks();

  // Bucket 1.
  folly::StringPiece key1 = "key1";
  // different keys are not affected.
  assertVecEq(tracker.recordAndPopulateAccessFeatures(key1), {0, 0, 0});
  // key1: {0, 1, 0}
  advanceTicks();

  assertVecEq(tracker.recordAndPopulateAccessFeatures(key1), {1, 0, 0});
  // key1: {0, 2, 0}
  advanceTicks();

  // Bucket 2.
  assertVecEq(tracker.recordAndPopulateAccessFeatures(key0), {0, 0, 2});
  // key0: {2, 0, 1}
  advanceTicks();

  assertVecEq(tracker.recordAndPopulateAccessFeatures(key0), {1, 0, 2});
  // key0: {2, 0, 2}
  advanceTicks();
  // Bucket 3.
  // Rotated, first 0 is discarded.
  assertVecEq(tracker.recordAndPopulateAccessFeatures(key0), {0, 2, 0});
  advanceTicks();
  // key0: {2, 0, 2, 1}
  assertVecEq(tracker.recordAndPopulateAccessFeatures(key0), {1, 2, 0});
  advanceTicks();
  // key0: {2, 0, 2, 2}

  // Bucket 4.
  // counts in bucket 1 was cleared and reset at last record.
  assertVecEq(tracker.recordAndPopulateAccessFeatures(key0), {0, 2, 2});
  advanceTicks();
  // key0: {2, 0, 2, 2, 1}
  // The other key's buckets were moved forward as well.
  assertVecEq(tracker.recordAndPopulateAccessFeatures(key1), {0, 0, 0});
  // key1: {0, 2, 0, 0, 1}
}

TEST_P(AccessTrackerTest, separateAccessTrackingAndPopulate) {
  auto config = AccessTracker::Config();
  config.numBuckets = 3;
  config.useCounts = GetParam();
  // Moves to a new bucket every two ticks.
  config.numTicksPerBucket = 2;
  config.ticker = ticker_;
  initializeTicks(
      config.numTicksPerBucket, config.numBuckets, config.numTicksPerBucket);
  auto tracker = AccessTracker(std::move(config));

  folly::StringPiece key = "key0";
  // Access is counted before recording.
  assertVecEq(tracker.getAccesses(key), {0, 0, 0});
  tracker.recordAccess(key);
  // Access is counted after recording access.
  assertVecEq(tracker.getAccesses(key), {1, 0, 0});

  tracker.recordAccess(key);
  tracker.recordAccess(key);
  // Multiple accesses counted together.
  assertVecEq(tracker.getAccesses(key), {3, 0, 0});

  tracker.recordAccess(key);
  advanceTicks();
  tracker.recordAccess(key);
  advanceTicks();
  tracker.recordAccess(key);
  // Multiple accesses across multiple buckets
  assertVecEq(tracker.getAccesses(key), {1, 5, 0});
}

} // namespace cachelib
} // namespace facebook
